"""
提示词工程工作流定义
支持人工介入反馈的提示词工程工作流实现
"""

from typing import Dict, Any
from langgraph.graph import StateGraph, END
from src.research_core.prompt_eng_state import PromptEngineeringState
from src.research_core.prompt_eng_agents import (
    analyze_prompt_requirement,
    design_initial_prompt,
    evaluate_prompt,
    optimize_prompt,
    process_human_feedback,
    finalize_prompt,
    should_continue_analysis,
    should_process_feedback,
    ab_test_prompts,
    collect_user_feedback,
    auto_optimize_prompt,
    test_prompt_in_production
)

def create_prompt_engineering_workflow_with_feedback():
    """
    创建支持人工介入反馈的提示词工程工作流
    包含完整的提示词迭代流程，支持人工反馈介入
    """
    # 创建状态图并添加节点
    workflow = StateGraph(PromptEngineeringState)
    
    # 添加所有工作流节点
    workflow.add_node("analyze_requirement", analyze_prompt_requirement)
    workflow.add_node("design_prompt", design_initial_prompt)
    workflow.add_node("check_feedback", process_human_feedback)
    workflow.add_node("evaluate_prompt", evaluate_prompt)
    workflow.add_node("optimize_prompt", optimize_prompt)
    workflow.add_node("finalize_prompt", finalize_prompt)

    # 设置入口点
    workflow.set_entry_point("analyze_requirement")

    # 建立节点连接关系
    workflow.add_edge("analyze_requirement", "design_prompt")
    workflow.add_edge("design_prompt", "check_feedback")
    
    # 添加条件边 - 检查是否有反馈需要处理
    workflow.add_conditional_edges(
        "check_feedback",
        should_process_feedback,
        {
            "process_feedback": "optimize_prompt",  # 如果有反馈，进入优化节点
            "evaluate": "evaluate_prompt"          # 如果没有反馈，直接进入评估
        }
    )

    # 添加条件边 - 评估后决定是否优化
    workflow.add_conditional_edges(
        "evaluate_prompt",
        should_continue_analysis,
        {
            "optimize": "optimize_prompt",  # 如果需要优化，继续优化
            "finalize": "finalize_prompt"   # 如果不需要优化，进入最终确定
        }
    )

    # 优化后的提示词需要重新检查反馈
    workflow.add_edge("optimize_prompt", "check_feedback")
    
    # 最终节点连接到结束
    workflow.add_edge("finalize_prompt", END)

    return workflow.compile()

def create_simple_prompt_workflow_with_feedback():
    """
    创建简化的提示词工作流（仅分析和设计），但仍支持反馈
    适用于快速创建和确认基础提示词的场景
    """
    workflow = StateGraph(PromptEngineeringState)
    
    # 添加节点
    workflow.add_node("analyze_requirement", analyze_prompt_requirement)
    workflow.add_node("design_prompt", design_initial_prompt)
    workflow.add_node("check_feedback", process_human_feedback)
    workflow.add_node("finalize_prompt", finalize_prompt)
    
    # 设置流程
    workflow.set_entry_point("analyze_requirement")
    workflow.add_edge("analyze_requirement", "design_prompt")
    workflow.add_edge("design_prompt", "check_feedback")
    workflow.add_edge("check_feedback", "finalize_prompt")
    workflow.add_edge("finalize_prompt", END)
    
    return workflow.compile()

def create_advanced_prompt_workflow():
    """
    创建高级提示词工程工作流，包含A/B测试、自动优化和生产环境测试
    适用于对提示词质量要求较高的场景
    """
    workflow = StateGraph(PromptEngineeringState)
    
    # 添加所有工作流节点
    workflow.add_node("analyze_requirement", analyze_prompt_requirement)
    workflow.add_node("design_prompt", design_initial_prompt)
    workflow.add_node("check_feedback", process_human_feedback)
    workflow.add_node("evaluate_prompt", evaluate_prompt)
    workflow.add_node("ab_test", ab_test_prompts)
    workflow.add_node("collect_feedback", collect_user_feedback)
    workflow.add_node("auto_optimize", auto_optimize_prompt)
    workflow.add_node("production_test", test_prompt_in_production)
    workflow.add_node("finalize_prompt", finalize_prompt)

    # 设置入口点
    workflow.set_entry_point("analyze_requirement")

    # 建立节点连接关系
    workflow.add_edge("analyze_requirement", "design_prompt")
    workflow.add_edge("design_prompt", "check_feedback")
    
    # 添加条件边 - 检查是否有反馈需要处理
    workflow.add_conditional_edges(
        "check_feedback",
        should_process_feedback,
        {
            "process_feedback": "auto_optimize",  # 如果有反馈，进入自动优化节点
            "evaluate": "evaluate_prompt"         # 如果没有反馈，直接进入评估
        }
    )

    # 评估后进行A/B测试
    workflow.add_edge("evaluate_prompt", "ab_test")
    
    # A/B测试后收集用户反馈
    workflow.add_edge("ab_test", "collect_feedback")
    
    # 收集反馈后进行自动优化
    workflow.add_edge("collect_feedback", "auto_optimize")
    
    # 自动优化后进行生产环境测试
    workflow.add_edge("auto_optimize", "production_test")
    
    # 生产环境测试后决定是否继续优化
    workflow.add_conditional_edges(
        "production_test",
        should_continue_analysis,
        {
            "optimize": "auto_optimize",     # 如果需要继续优化，回到自动优化
            "finalize": "finalize_prompt"    # 如果不需要优化，进入最终确定
        }
    )

    # 最终节点连接到结束
    workflow.add_edge("finalize_prompt", END)

    return workflow.compile()

def create_multi_version_prompt_workflow():
    """
    创建多版本并行优化的提示词工程工作流
    同时生成多个优化版本，通过A/B测试选择最佳版本
    """
    workflow = StateGraph(PromptEngineeringState)
    
    # 添加所有工作流节点
    workflow.add_node("analyze_requirement", analyze_prompt_requirement)
    workflow.add_node("design_prompt", design_initial_prompt)
    workflow.add_node("check_feedback", process_human_feedback)
    workflow.add_node("evaluate_prompt", evaluate_prompt)
    workflow.add_node("multi_version_optimize", _multi_version_optimize)
    workflow.add_node("ab_test_versions", _ab_test_prompt_versions)
    workflow.add_node("select_best_version", _select_best_prompt_version)
    workflow.add_node("finalize_prompt", finalize_prompt)

    # 设置入口点
    workflow.set_entry_point("analyze_requirement")

    # 建立节点连接关系
    workflow.add_edge("analyze_requirement", "design_prompt")
    workflow.add_edge("design_prompt", "check_feedback")
    
    # 添加条件边 - 检查是否有反馈需要处理
    workflow.add_conditional_edges(
        "check_feedback",
        should_process_feedback,
        {
            "process_feedback": "multi_version_optimize",  # 如果有反馈，进入多版本优化节点
            "evaluate": "evaluate_prompt"                  # 如果没有反馈，直接进入评估
        }
    )

    # 评估后进行多版本优化
    workflow.add_edge("evaluate_prompt", "multi_version_optimize")
    
    # 多版本优化后进行A/B测试
    workflow.add_edge("multi_version_optimize", "ab_test_versions")
    
    # A/B测试后选择最佳版本
    workflow.add_edge("ab_test_versions", "select_best_version")
    
    # 选择最佳版本后决定是否继续优化
    workflow.add_conditional_edges(
        "select_best_version",
        should_continue_analysis,
        {
            "optimize": "multi_version_optimize",  # 如果需要继续优化，回到多版本优化
            "finalize": "finalize_prompt"          # 如果不需要优化，进入最终确定
        }
    )

    # 最终节点连接到结束
    workflow.add_edge("finalize_prompt", END)

    return workflow.compile()

def _multi_version_optimize(state: PromptEngineeringState) -> Dict[str, Any]:
    """
    多版本并行优化提示词
    """
    # 延迟初始化组件
    from src.research_core.prompt_eng_agents import _initialize_components
    _initialize_components()
    
    # 获取当前提示词
    current_prompt = state.get('current_prompt', '')
    requirement = state.get('requirement', '')
    
    # 调用LLM生成多个优化版本
    # 这里应该实现调用LLM生成多个版本的逻辑
    # 简化实现，实际应该调用LLM生成多个版本
    import random
    prompt_versions = {}
    version_names = ["version_a", "version_b", "version_c"]
    
    for i, version_name in enumerate(version_names):
        # 为每个版本添加不同的优化策略
        optimization_strategies = [
            f"优化策略A：增强清晰度和结构化",
            f"优化策略B：提高具体性和完整性",
            f"优化策略C：加强相关性和适应性"
        ]
        
        prompt_versions[version_name] = f"{current_prompt}\n\n优化说明: {optimization_strategies[i]}"
    
    return {
        "prompt_versions": prompt_versions,
        "current_stage": "multi_version_optimize"
    }

def _ab_test_prompt_versions(state: PromptEngineeringState) -> Dict[str, Any]:
    """
    对多个提示词版本进行A/B测试
    """
    # 获取提示词版本
    prompt_versions = state.get('prompt_versions', {})
    
    # 模拟真实的A/B测试逻辑
    ab_test_results = {}
    
    # 为每个版本生成测试结果
    for version, prompt in prompt_versions.items():
        # 模拟质量评分（基于更复杂的评估逻辑）
        from src.research_core.prompt_eng_agents import advanced_evaluate_prompt_quality
        evaluation_result = advanced_evaluate_prompt_quality(prompt, state.get('requirement', ''))
        overall_score = evaluation_result.get('overall_score', 0.0)
        
        ab_test_results[version] = {
            "prompt": prompt,
            "score": overall_score,
            "details": evaluation_result,
            "feedback": "测试完成"
        }
    
    return {
        "ab_test_results": ab_test_results,
        "current_stage": "ab_test_versions"
    }

def _select_best_prompt_version(state: PromptEngineeringState) -> Dict[str, Any]:
    """
    选择最佳提示词版本
    """
    ab_test_results = state.get('ab_test_results', {})
    
    if not ab_test_results:
        return {
            "current_stage": "select_best_version"
        }
    
    # 选择得分最高的版本
    best_version = max(ab_test_results.items(), key=lambda x: x[1]['score'])
    
    # 记录选择原因
    selection_reasoning = f"选择版本 {best_version[0]}，得分 {best_version[1]['score']:.2f}，优于其他版本"
    
    return {
        "current_prompt": best_version[1]['prompt'],
        "quality_score": best_version[1]['score'],
        "quality_evaluation_details": best_version[1]['details'],
        "selection_reasoning": selection_reasoning,
        "current_stage": "select_best_version"
    }

# 按照LangGraph Studio项目导出规范，导出graph变量
graph = create_prompt_engineering_workflow_with_feedback()